package main

import (
	"context"
	"fmt"
	"time"

	"go.uber.org/zap"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

type GormLogger struct {
	SlowThreshold time.Duration
	logger        *zap.Logger
}

func NewGormLogger(logger *zap.Logger) *GormLogger {
	return &GormLogger{
		//SlowThreshold: 2 * time.Millisecond,
		logger: logger,
	}
}

var _ logger.Interface = (*GormLogger)(nil)

func (l *GormLogger) LogMode(lev logger.LogLevel) logger.Interface {
	return &GormLogger{}
}

func (l *GormLogger) Info(ctx context.Context, msg string, data ...interface{}) {
	l.logger.Info(msg, zap.Any("data", data))
}

func (l *GormLogger) Warn(ctx context.Context, msg string, data ...interface{}) {
	l.logger.Warn(msg, zap.Any("data", data))
}

func (l *GormLogger) Error(ctx context.Context, msg string, data ...interface{}) {
	l.logger.Error(msg, zap.Any("data", data))
}

func (l *GormLogger) Trace(ctx context.Context, begin time.Time, fc func() (sql string, rowsAffected int64), err error) {
	elapsed := time.Since(begin)
	sql, rows := fc()
	//if err != nil {
	//	l.logger.Error("SQL ERROR",
	//		zap.String("biz", fmt.Sprintf("%v", ctx.Value("biz"))),
	//		zap.String("sql", sql),
	//		zap.Int64("rows", rows),
	//		zap.Duration("elapsed", elapsed),
	//		zap.Error(err),
	//	)
	//}
	//if l.SlowThreshold != 0 && elapsed > l.SlowThreshold {
	//	l.logger.Warn("Database Slow Log",
	//		zap.String("biz", fmt.Sprintf("%v", ctx.Value("biz"))),
	//		zap.String("sql", sql),
	//		zap.Int64("rows", rows),
	//		zap.Duration("elapsed", elapsed),
	//	)
	//}

	l.logger.Info("Database Info Log",
		zap.String("biz", fmt.Sprintf("%v", ctx.Value("biz"))),
		zap.String("sql", sql),
		zap.Int64("rows", rows),
		zap.Duration("elapsed", elapsed),
	)

}

func main() {
	logger, _ := zap.NewProduction()
	defer logger.Sync()
	db, err := gorm.Open(mysql.Open(`root:12345678@tcp(127.0.0.1:3306)/test1`), &gorm.Config{
		QueryFields: true,
		Logger:      NewGormLogger(logger),
	})
	if err != nil {
		panic(err)
	}

	var count int
	todayBatchTotalCountCtx := context.WithValue(context.Background(), "biz", "我的业务")
	err = db.WithContext(todayBatchTotalCountCtx).Table("test").
		Select("count(1)").Row().Scan(&count)
	if err != nil {
		fmt.Println(err)
	}
}

/*  创建测试数据
CREATE TABLE test1.test (
    id INT AUTO_INCREMENT PRIMARY KEY,
    name VARCHAR(50) NOT NULL,
    age INT,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
*/
